package vectorindex_v2

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"strconv"
	"sync"

	"github.com/cloudwego/eino/schema"
)

// FaissRetrieverConfig 包含 Faiss 检索器的配置
type FaissRetrieverConfig struct {
	// 索引文件路径
	IndexPath string
	DocPath   string

	// 向量维度
	VectorDim int

	// 默认检索参数
	DefaultTopK int
}

// WithFaissRetrieverConfig 创建 Faiss 检索器配置选项
func WithFaissRetrieverConfig(config *FaissRetrieverConfig) RetrieveOption {
	return func(opt *RetrieveOptions) {
		// 这里简化处理，实际项目中可能需要更复杂的选项传递机制
	}
}

// FaissRetriever 实现 VectorRetriever 接口
type FaissRetriever struct {
	index      *faiss.IndexImpl
	config     *FaissRetrieverConfig
	docs       map[int64]*schema.Document
	docsMutex  sync.RWMutex
	indexMutex sync.RWMutex
}

// NewFaissRetriever 创建新的 Faiss 检索器
func NewFaissRetriever(ctx context.Context, config *FaissRetrieverConfig) (VectorRetriever, error) {
	// 验证必需配置
	if config.IndexPath == "" {
		return nil, fmt.Errorf("index path is required")
	}
	if config.DocPath == "" {
		return nil, fmt.Errorf("document path is required")
	}
	if config.VectorDim <= 0 {
		return nil, fmt.Errorf("vector dimension must be positive")
	}

	// 设置默认值
	if config.DefaultTopK <= 0 {
		config.DefaultTopK = 10
	}

	// 加载索引
	idx, err := faiss.ReadIndex(config.IndexPath, faiss.IO_FLAG_READ_ONLY)
	if err != nil {
		return nil, fmt.Errorf("failed to read index from file: %w", err)
	}

	// 检查索引维度是否匹配
	dim := idx.D()
	if dim != config.VectorDim {
		return nil, fmt.Errorf("index dimension mismatch: expected %d, got %d", config.VectorDim, dim)
	}

	// 加载文档
	docs := make(map[int64]*schema.Document)
	data, err := os.ReadFile(config.DocPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read document file: %w", err)
	}

	if err := json.Unmarshal(data, &docs); err != nil {
		return nil, fmt.Errorf("failed to unmarshal documents: %w", err)
	}

	return &FaissRetriever{
		index:  idx,
		config: config,
		docs:   docs,
	}, nil
}

// Retrieve 实现 retriever.Retriever 接口的 Retrieve 方法
func (fr *FaissRetriever) Retrieve(ctx context.Context, query string, opts ...RetrieveOption) ([]*schema.Document, error) {
	// 解析选项
	commonOpts := GetRetrieveCommonOptions(nil, opts...)

	// 确定 TopK
	topK := fr.config.DefaultTopK
	if commonOpts.TopK != nil && *commonOpts.TopK > 0 {
		topK = *commonOpts.TopK
	}

	// 检查是否提供了嵌入器
	if commonOpts.Embedding == nil {
		return nil, fmt.Errorf("embedding component is required for retrieval")
	}

	// 生成查询向量
	vectors, err := commonOpts.Embedding.EmbedStrings(ctx, []string{query})
	if err != nil {
		return nil, fmt.Errorf("failed to embed query: %w", err)
	}
	if len(vectors) == 0 || len(vectors[0]) == 0 {
		return nil, fmt.Errorf("failed to generate valid embedding for query")
	}

	// 转换为 float32
	queryVector := make([]float32, len(vectors[0]))
	for i, v := range vectors[0] {
		queryVector[i] = float32(v)
	}

	// 执行检索
	fr.indexMutex.RLock()
	k := topK
	distances := make([]float32, k)
	labels := make([]int64, k)
	err = fr.index.Search(1, queryVector, k, distances, labels)
	fr.indexMutex.RUnlock()

	if err != nil {
		return nil, fmt.Errorf("failed to search index: %w", err)
	}

	// 构建结果文档
	var documents []*schema.Document
	fr.docsMutex.RLock()
	defer fr.docsMutex.RUnlock()

	for i := 0; i < k; i++ {
		// 检查是否找到有效结果
		if labels[i] == -1 {
			continue
		}

		doc, exists := fr.docs[labels[i]]
		if !exists {
			continue
		}

		// 创建文档副本并添加相似度分数
		resultDoc := &schema.Document{
			ID:       strconv.FormatInt(labels[i], 10),
			Content:  doc.Content,
			MetaData: make(map[string]interface{}),
		}

		// 复制元数据
		for key, value := range doc.MetaData {
			resultDoc.MetaData[key] = value
		}

		// 添加相似度分数（对于L2距离，我们取负数作为相似度分数）
		resultDoc.MetaData["similarity_score"] = -distances[i]

		documents = append(documents, resultDoc)
	}

	return documents, nil
}

// Close 关闭检索器并释放资源
func (fr *FaissRetriever) Close() error {
	// Faiss-Go 的索引不需要显式销毁
	fr.index = nil
	return nil
}
